#!/usr/bin/env python
"""
This package selects a random  set of points from an image and writes them to
disk either in an XML scheme that is compatiable with the MITK mps data structure
or as a cPickle file wiwth my CritPoint structure.

Points are selected from regions that satisfy simple threshold criteria specified by the user. Optional morphological operators can be applied to modify the 
regions prior to sampling
"""
import itk
import imageTools.ITKUtils.io as io
import imageTools.ITKUtils.query as query
import imageTools.ITKUtils.ifilter as ifilter
import imageTools.ITKUtils.view as view
import imageTools.browse.CritPoint as CritPoint
import numpy as na
import numpy.random as ra
import MyIU
import sys
import imageTools.ITKUtils.itkmath as itkmath
import os
import cPickle

from optparse import OptionParser

def getParser():
    """create an OptionParser instance and create parsing rules"""
    try:
        parser = OptionParser()
        parser.add_option("-t", "--thresh", dest='threshold',nargs=2, default=(-150, 230), type='int')
        parser.add_option("-k", "--kernel", dest='kernel',nargs=3,
                          default=(1,1,1),type='int')
        parser.add_option("-n", "--numPoints", dest='numPoints',default=200,type='int')
        parser.add_option("-r", "--rescale",action='store_true', dest='rescale',default=False)
        parser.add_option("-o", "--output", dest="outName", default="sample.mps")
        parser.add_option("-p", "--points", dest="pointsFile", default="")
        parser.add_option("-d", "--dilate",action='store_true', dest='dilate',default=True)
        parser.add_option("-v", "--view",action='store_true',
                          dest='view',default=False)
        return parser
    except Exception,  error:
        print "failed to generate parser",  error

def dilateMask(img, kernel= None):
    """dilate the binary mask passed in with the positional argument img.
    A default kernel size of (2,2,1) is assumed. 

    Prior to dilation an itk.BinaryMedianImageFilter is used with a (1,1,1)
    kerenl
    
    Dilation is applied using the itk.BinaryDilateImageFilter"""
    try:
        # If user didn't provide a kernel size, use the default value
        if( kernel == None ):
            kernel = (2,2,1)
        median = itk.BinaryMedianImageFilter.IUC3IUC3.New()
        median.SetRadius([1,1,1])
        
        median.SetForegroundValue(1)
        median.SetInput(img)
        dilate = itk.BinaryDilateImageFilter.IUC3IUC3SE3.New()
        dilate.SetForegroundValue(1)
        kd = dilate.GetKernel()
        kd.SetRadius(kernel)
        dilate.SetKernel(kd)
        print "Get first median image"
        img2 = median.GetOutput()
        dilate.SetInput(img2)
        print "dilate image"
        # this is a really, really slow step
        img3 = dilate.GetOutput()
        img3.Update()
        #diffImg = itkmath.subtractImages(img3,img2)
        #print itk.range(diffImg)
        #view.view3DImage([img,img2,img3]) 
        return img3
    except Exception,  error:
        print "failed in dilateMask",  error
def orImages(m1,m2):        
    if( m2 == None ):
        return m1
    orFilter = itk.OrImageFilter.IUC3IUC3IUC3.New()
    orFilter.SetInput1(m1)
    orFilter.SetInput2(m2)
    orFilter.Update()
    return orFilter.GetOutput()
       
def iterativelyGetMask(img,  
                       rescale=True,  
                       thresh=(-150, 230),
                       interactive=False):
    """
    Takes an image and a tuple of threshold and create a binary mask defined 
    by the thresholds. Present the image to the user to accept the mask or 
    modify the thresholds. Currently this relies on a hokey  display solution
    
    img: an itk Imge currently assumes image type is itk.Image[itk.F,3]
    rescale: Binary flag. If true image is rescaled to (0,1000)
    interactive: Binary flag. If true, the user is asked to veify intermediate
    images prior to accepting the segmentation.
    
    """
    try:
        iType = itk.Image[itk.F, 3] 
        if(rescale):
            img = ifilter.rescaleImage(img, [0, 100])
        while(True):
            mask = ifilter.thresholdImage(img, thresh, (0, 1), iType=iType) # I need to change this so I don't need to specify iType
            img2 = itkmath.multiplyImages(img, mask)
            if( not interactive ):
                break
            view.view3DImage([img, img2], title="Threshold=%d,%ds"%thresh)
            q = raw_input("Hit enter to accept current segmentation or type in a new theshold pair as low, high: ")
            if( not q):
                break
            else:
                try:
                    q = q.split(',')
                    thresh = int(q[0]), int(q[1])
                    print "recomputing mask with thresholds %d and %d"%(thresh[0], thresh[1])
                except:
                    print "invalid entry. Repeating with current threshold"
        print "completed mask generation"           
        return mask, thresh
    except Exception,  error:
        print "failed in iterativelyGetMask()",  error


def samplePoints(img,  numPoints=20):
    """sample critical points from the user suppied image.
    
        img: an itk img provided by the user
        numPoints: number of random points to sample """
    try:

        mask = itk.ImageMaskSpatialObject._3.New()
        mask.SetImage(img)
        mask.Update()
        array = itk.PyBuffer.IUC3.GetArrayFromImage(img)
        indices = na.where(array. flat==1)[0]
        rinds = ra.randint(0,len(indices),numPoints)
        samples = indices.take(rinds)
        crds = MyIU.get_crds(samples,array.shape)
        p = itk.Point.D3()
        iwt = mask.GetIndexToObjectTransform()
        points = []
        for i in range(len(crds[0])):
            print i
            p.SetElement(0, crds[0][i])
            p.SetElement(1, crds[1][i])
            p.SetElement(2, crds[2][i])
            p2 = iwt.TransformPoint(p)
            query.print_itkPoint(p, 'Index')
            query.print_itkPoint(p2, 'World Coordinate')
            a = query.get_itkPoint(p)
            a2 = query.get_itkPoint(p2)
            cp = CritPoint.CritPoint()
            cp.extend(a)
            cp.setLabel("unlabeled")
            cp.worldCoordinates = a2
            wc = CritPoint.CritPoint()
            wc.extend(a2)
            cp.setWorldCoordinates(wc)
            points.append(cp)
        return points
    except Exception,  error:
        print "failed in samplePoints",  error

def readPoints(fname):
    try:
        fo = open(fname,'rb')
        points = cPickle.load(fo)
        return points
    except Exception, error:
        print "failed in readPoints", error

def getUniquePoints(p1,p2):
    """Given two lists (p1 and p2) of critical points, return a single
    list of unique points between the two lists"""
    try:
        for p in p2:
            if p not in p1:
                p1.append(p)
        return p1
    except Exception, error:
        print "faile din getUniquePoints", error


def main():
    try:
        
        parser = getParser()
        (options, args) = parser.parse_args()
        iType = itk.Image[itk.F, 3]
        oType = itk.Image[itk.UC, 3]
        finalMask = None
        for file in args:
            print file
            img = io.readImage(str(file),iType=iType)
            mask,  thresh = iterativelyGetMask(img,  rescale=options.rescale,  thresh=options.threshold)
            if( options.view ):
                view.view3DImage([mask])
            if(options.dilate):
                print "modifying mask"
                mask=dilateMask(mask,kernel=options.kernel)
                if( options.view ):
                    view.view3DImage([mask])
            finalMask = orImages(mask,finalMask)
        
        print "sampling points"
        newPoints = samplePoints(finalMask,  numPoints=options.numPoints)
        if( options.pointsFile ):
            oldPoints = readPoints(options.pointsFile)
            newPoints = getUniquePoints(newPoints,oldPoints)
            fo = open(options.pointsFile,'wb')
            cPickle.dump(newPoints,fo)
        else:
            fo = open(options.outName+"Sampled.pckle","wb")
            cPickle.dump(newPoints,fo)

        io.writeImage(finalMask,options.outName+"SampledMask.mha")
        fo = open(options.outName+"Sampled.info", "w")
        fo.write("%s\n"%sys.argv.__str__())
        fo.write("%s\n"%options.__str__())
        fo.write("thresh:%s\n"%thresh.__str__())
        fo.close()
    except Exception,  error:
        print "failed in main",  error
        
if __name__=='__main__':
    main()
        
        
        
